import subprocess
import logging
import os
from typing import List, Optional

logger = logging.getLogger(__name__)

class NetworkIsolation:
    """网络隔离工具"""
    
    def __init__(self):
        self.isolated_pids: List[int] = []
        self.chain_name = "SANDBOX_BLOCK"
        self._setup_iptables_chain()
    
    def _setup_iptables_chain(self):
        """设置iptables链用于沙箱隔离"""
        try:
            # 创建自定义链（如果不存在）
            subprocess.run([
                "iptables", "-t", "filter", "-N", self.chain_name
            ], check=False, capture_output=True)
            
            # 在OUTPUT链中添加跳转到自定义链的规则
            subprocess.run([
                "iptables", "-t", "filter", "-C", "OUTPUT", "-j", self.chain_name
            ], check=False, capture_output=True)
            
            if subprocess.run([
                "iptables", "-t", "filter", "-C", "OUTPUT", "-j", self.chain_name
            ], capture_output=True).returncode != 0:
                subprocess.run([
                    "iptables", "-t", "filter", "-I", "OUTPUT", "1", "-j", self.chain_name
                ], check=True, capture_output=True)
            
            logger.debug(f"iptables链 {self.chain_name} 设置完成")
            
        except subprocess.CalledProcessError as e:
            logger.warning(f"设置iptables链失败（可能需要root权限）: {e}")
        except FileNotFoundError:
            logger.warning("iptables未找到，网络隔离将被禁用")
    
    def isolate_process(self, pid: int) -> bool:
        """隔离指定进程的网络访问"""
        try:
            # 阻止指定PID的所有出站连接
            subprocess.run([
                "iptables", "-t", "filter", "-A", self.chain_name,
                "-m", "owner", "--pid-owner", str(pid),
                "-j", "DROP"
            ], check=True, capture_output=True)
            
            self.isolated_pids.append(pid)
            logger.debug(f"已隔离进程 {pid} 的网络访问")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.warning(f"隔离进程 {pid} 网络访问失败: {e}")
            return False
        except FileNotFoundError:
            logger.warning("iptables未找到，无法隔离网络")
            return False
    
    def remove_process_isolation(self, pid: int) -> bool:
        """移除指定进程的网络隔离"""
        try:
            subprocess.run([
                "iptables", "-t", "filter", "-D", self.chain_name,
                "-m", "owner", "--pid-owner", str(pid),
                "-j", "DROP"
            ], check=True, capture_output=True)
            
            if pid in self.isolated_pids:
                self.isolated_pids.remove(pid)
            
            logger.debug(f"已移除进程 {pid} 的网络隔离")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.warning(f"移除进程 {pid} 网络隔离失败: {e}")
            return False
    
    def cleanup_all(self):
        """清理所有网络隔离规则"""
        try:
            # 移除所有隔离的进程
            for pid in self.isolated_pids.copy():
                self.remove_process_isolation(pid)
            
            # 清空自定义链
            subprocess.run([
                "iptables", "-t", "filter", "-F", self.chain_name
            ], check=False, capture_output=True)
            
            logger.debug("已清理所有网络隔离规则")
            
        except Exception as e:
            logger.warning(f"清理网络隔离规则失败: {e}")

# 全局网络隔离实例
network_isolation = NetworkIsolation()

def disable_network_for_process(pid: int) -> bool:
    """为进程禁用网络访问"""
    return network_isolation.isolate_process(pid)

def enable_network_for_process(pid: int) -> bool:
    """为进程恢复网络访问"""
    return network_isolation.remove_process_isolation(pid)

def cleanup_network_isolation():
    """清理所有网络隔离"""
    network_isolation.cleanup_all() 